package core

import (
	"context"
	"encoding/binary"
	"fmt"
	"github.com/kataras/golog"
	"math/rand"
	"net"
	"ntrace/enviroments"
	"ntrace/utils/routines"
	"time"
)

type UDPSender struct {
	BaseSender
}

func NewUDPSender(ctx context.Context, task *TraceTask, stat *TraceStat) *UDPSender {
	return &UDPSender{
		BaseSender: BaseSender{
			Ctx:  ctx,
			Name: "ICMPSender",
			Task: task,
			Stat: stat,
		},
	}
}

func (s *UDPSender) StartSender() error {
	routines.GoWithRecover("UDPSender", s.doSend)
	return nil
}

func (s *UDPSender) doSend() {
	rc, err := OpenSenderIP4RawPacketConn("ip4:udp", "0.0.0.0")
	if err != nil {
		return
	}
	s.Conn = rc

	task := s.Task
	payload := AllocTxBuf(task.MsgLen)
	for turn := 0; turn < task.Count; turn++ {
		for srcPort := task.SrcPortMin; srcPort <= task.SrcPortMax; srcPort++ {
			for dstPort := task.DstPortMin; dstPort <= task.DstPortMax; dstPort++ {
				ipID := uint16(rand.Intn(1 << 16))
				for ttl := uint8(1); ttl <= task.MaxTTL; ttl++ {
					select {
					case <-s.Ctx.Done():
						s.StopSender()
						return
					default:
						ipID++
						_ = s.sendOnePkt(turn, task.SrcIPs[0], ipID, srcPort, dstPort, ttl, payload)
					}
				}

				time.Sleep(time.Duration(task.FlowIntervalMs) * time.Millisecond)
			}
		}
	}

	// 发送结束
	s.StopSender()
	TrigExit(s.Task.TimeoutMs)
}

func (s *UDPSender) sendOnePkt(turn int, srcIP string, ipID uint16, srcPort uint16, dstPort uint16, ttl uint8, payload []byte) error {
	task := s.Task
	target := task.DstIP
	pid := enviroments.PID

	golog.Debugf("send, id: %d, seq: %d, srcIP: %s, dstIP: %s, srcPort: %d, dstPort: %d",
		pid, ttl, srcIP, target, srcPort, dstPort)

	txTs := time.Now().UnixNano()
	binary.BigEndian.PutUint16(payload[8:10], uint16(pid))
	binary.BigEndian.PutUint16(payload[10:12], ipID)
	binary.BigEndian.PutUint64(payload[12:20], uint64(txTs))
	msg, err := EncodeIPUDPPacket(srcIP, target, ipID, task.TOS, ttl, srcPort, dstPort, payload)
	if err != nil {
		return err
	}
	err = s.Conn.SetWriteDeadline(time.Now().Add(500 * time.Millisecond))
	if err != nil {
		return err
	}
	_, err = s.Conn.WriteToIP(msg, &net.IPAddr{IP: net.ParseIP(target)})
	if err != nil {
		return err
	}

	s.record(turn, srcIP, ttl, ipID, srcPort, dstPort, txTs)
	return nil
}

func (s *UDPSender) record(turn int, srcIP string, ttl uint8, ipID uint16, srcPort uint16, dstPort uint16, txTs int64) {
	key := fmt.Sprintf("%s-%d-%d-%d", s.Task.DstIP, ipID, srcPort, dstPort)
	golog.Debugf("sender ttl: %d, key: %s", ttl, key)
	pktRecord := &PktRecord{
		Key:     key,
		Turn:    turn,
		SrcIP:   srcIP,
		DstIP:   s.Task.DstIP,
		SrcPort: srcPort,
		DstPort: dstPort,
		TTL:     ttl,

		TxTs: txTs,
	}

	s.Stat.PktRecordBucket.Add(pktRecord)
}
